"""
AI对话上下文管理器
负责管理对话的上下文记忆，包括Token计算、上下文压缩等功能
"""

import json
from typing import List, Dict, Any
from django.core.cache import cache
from django.utils import timezone
from .models import AIConversation, AIMessage, AIRole


class ConversationContextManager:
    """对话上下文管理器"""
    
    def __init__(self, conversation: AIConversation):
        self.conversation = conversation
        self.cache_key = f"conversation_context:{conversation.id}"
        self.cache_timeout = 1800  # 30分钟
    
    def get_context_messages(
        self,
        message_limit: int = 10,
        token_limit: int = 120000,
        include_system_prompt: bool = True
    ) -> Dict[str, Any]:
        """
        获取对话上下文消息 - 使用滑动窗口实现

        Args:
            message_limit: 消息数量限制
            token_limit: Token数量限制
            include_system_prompt: 是否包含系统提示词

        Returns:
            包含上下文信息的字典
        """
        # 尝试从缓存获取
        cache_key = f"{self.cache_key}:{message_limit}:{token_limit}:{include_system_prompt}"
        cached_context = cache.get(cache_key)
        if cached_context:
            return cached_context

        context_messages = []
        total_tokens = 0
        system_prompt = None

        # 1. 添加系统提示词（如果需要）
        if include_system_prompt and self.conversation.ai_role:
            system_prompt = self.conversation.ai_role.system_prompt
            system_tokens = self._estimate_tokens(system_prompt)
            total_tokens += system_tokens

            context_messages.append({
                'role': 'system',
                'content': system_prompt,
                'token_count': system_tokens,
                'created_at': self.conversation.created_at,
                'is_compressed': False
            })

        # 2. 获取最近的消息（滑动窗口）
        recent_messages = self.conversation.messages.order_by('-created_at')

        # 3. 使用滑动窗口选择消息
        selected_messages = []
        for message in recent_messages:
            message_tokens = message.token_count or self._estimate_tokens(message.content)

            # 检查是否超过Token限制
            if total_tokens + message_tokens > token_limit:
                break

            # 检查是否超过消息数量限制
            if len(selected_messages) >= message_limit:
                break

            selected_messages.append(message)
            total_tokens += message_tokens

        # 4. 按时间正序排列消息（最早的在前）
        selected_messages.reverse()

        # 5. 转换为上下文格式
        for message in selected_messages:
            role = 'assistant' if message.ai_role else 'user'
            message_tokens = message.token_count or self._estimate_tokens(message.content)

            context_messages.append({
                'role': role,
                'content': message.content,
                'token_count': message_tokens,
                'created_at': message.created_at,
                'is_compressed': False
            })

        result = {
            'conversation_id': self.conversation.id,
            'total_messages': self.conversation.messages.count(),
            'context_messages': context_messages,
            'total_tokens': total_tokens,
            'system_prompt': system_prompt,
            'context_summary': None,
            'window_size': len(selected_messages)
        }

        # 缓存结果
        cache.set(cache_key, result, self.cache_timeout)

        return result
    
    def _estimate_tokens(self, text: str) -> int:
        """
        估算文本的Token数量
        简单估算：中文按字符数，英文按单词数*1.3
        """
        if not text:
            return 0
        
        # 简单的Token估算算法
        chinese_chars = len([c for c in text if '\u4e00' <= c <= '\u9fff'])
        other_chars = len(text) - chinese_chars
        
        # 中文字符大约1个字符=1个Token，英文大约4个字符=1个Token
        estimated_tokens = chinese_chars + (other_chars // 4)
        
        return max(1, estimated_tokens)
    

    
    def add_message_to_context(self, message: AIMessage):
        """添加新消息到上下文并更新缓存"""
        # 清除相关缓存（由于Django cache不支持模式删除，我们清除整个会话的缓存）
        self.clear_context_cache()
        
        # 更新会话的最后消息时间和Token使用量
        self.conversation.last_message_at = message.created_at
        if message.token_count:
            self.conversation.used_tokens += message.token_count
        self.conversation.save(update_fields=['last_message_at', 'used_tokens'])
    
    def clear_context_cache(self):
        """清除上下文缓存"""
        # 由于Django cache不支持模式删除，我们需要删除可能的缓存键
        # 尝试删除常用的参数组合
        common_params = [
            (10, 120000, True),
            (10, 120000, False),
            (20, 120000, True),
            (20, 120000, False),
            (5, 120000, True),
            (5, 120000, False),
        ]
        
        for message_limit, token_limit, include_system_prompt in common_params:
            cache_key = f"{self.cache_key}:{message_limit}:{token_limit}:{include_system_prompt}"
            cache.delete(cache_key)
        
        # 也删除基础缓存键
        cache.delete(self.cache_key)
    
    def get_context_for_ai_request(
        self,
        message_limit: int = 10,
        token_limit: int = 120000
    ) -> List[Dict[str, str]]:
        """
        获取用于AI请求的上下文格式

        Returns:
            符合OpenAI/Anthropic API格式的消息列表
        """
        context_data = self.get_context_messages(
            message_limit=message_limit,
            token_limit=token_limit,
            include_system_prompt=True
        )
        
        # 转换为AI API格式
        api_messages = []
        for msg in context_data['context_messages']:
            api_messages.append({
                'role': msg['role'],
                'content': msg['content']
            })
        
        return api_messages
